package com.lds.memory.kafka.partition;

import com.lds.memory.kafka.broker.status.Status;
import com.lds.memory.kafka.consumer.parallel.ConsumerTask;
import com.lds.memory.kafka.partition.reject.DefaultRejectedStrategy;
import com.lds.memory.kafka.partition.reject.RejectedStrategy;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

public class Partition<T> implements Storage<T>, Status {
    /**
     * Partition运行开关
     */
    private volatile boolean running = false;
    private ConsumerTask consumerTask;
    private final BlockingQueue<T> queue;
    private final RejectedStrategy rejectedStrategy;

    public Partition(int queueSize) {
        this(queueSize, null);
    }

    public Partition(int queueSize, RejectedStrategy rejectedStrategy) {
        this.queue = new ArrayBlockingQueue<>(queueSize);
        this.rejectedStrategy = rejectedStrategy == null ? new DefaultRejectedStrategy() : rejectedStrategy;
    }

    public void bind(ConsumerTask consumerTask) {
        this.consumerTask = consumerTask;
    }

    @Override
    public void store(T data) {
        if (consumerTask == null) return;
        if (!running) return;
        if (!queue.offer(data)) {
            rejectedStrategy.rejected(data);
        }
    }

    @Override
    public Collection<T> acquire() {
        if (consumerTask == null) return null;
        List<T> list = new ArrayList<>();
        queue.drainTo(list);
        return list;
    }

    @Override
    public void start() {
        if (consumerTask == null) return;
        if (running) return;
        running = true;
    }

    /**
     * 不再接收新数据
     */
    @Override
    public void pause() {
        if (consumerTask == null) return;
        if (!running) return;
        running = false;
    }

    @Override
    public void cancelPause() {
        if (consumerTask == null) return;
        if (running) return;
        running = true;
    }

    @Override
    public void shutdown() {
        pause();

        consumerTask.shutdown();
        // 保证尽快被垃圾回收
        consumerTask = null;
    }

}
